﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace Jarvis.Learners.DecisionTree
{
    class TreeNode
    {
        public String feature_name = "";
        public Dictionary<Object, TreeNode> children;
        public String label;
        public Double split = Double.PositiveInfinity;
        public object subitem = null;
        public double examples_count = 0;
        public TreeNode unknown_value;

        public TreeNode()
        {
            children = new Dictionary<object, TreeNode>();
        }

        public String GetLabel(Dictionary<String, Object> state)
        {
            if (label != null)
                return label;
            else
            {
                if (state.ContainsKey(feature_name))
                {
                    Object value = state[feature_name];

                    TreeNode node = null;

                    if (value is Double)
                    {
                        if ((Double) value >= split)
                            node = children[">="];
                        else
                            node = children["<"];
                    }
                    else if (value is List<object>)
                    {
                        if (((List<object>)value).Contains(subitem))
                            node = children["Contains"];
                        else if (children.ContainsKey("Does not contain"))
                            node = children["Does not contain"];
                    }
                    else if (children.ContainsKey(value))
                        node = children[value];

                    if (node != null)
                        return node.GetLabel(state);
                }
                else
                {
                    Double max_count = Double.NegativeInfinity;
                    String max_label = null;

                    // TODO: Change from just fetching the first node...

                    foreach (TreeNode node in children.Values)
                    {
                        String l = node.GetLabel(state);

                        if (l != null)
                        {
                            if (node.examples_count > max_count)
                            {
                                max_count = node.examples_count;
                                max_label = l;
                            }
                        }
                    }

                    if (unknown_value != null)
                    {
                        String l = unknown_value.GetLabel(state);

                        if (l != null)
                        {
                            max_count = unknown_value.examples_count;
                            max_label = l;
                        }
                    }

                    return max_label;
                }
            }

            return null;
        }

        protected String ToString(String indent)
        {
            if (label != null)
                return indent + "Prediction: " + label + Environment.NewLine;
            else
            {
                String s = indent + "Sensor: " + feature_name + Environment.NewLine;


                if (split != Double.PositiveInfinity)
                    s = s + indent + "Split Value: " + split + Environment.NewLine;

                if (subitem != null)
                    s = s + indent + "Contained Item: " + subitem + Environment.NewLine;

                s = s + indent + "Example Count: " + examples_count + Environment.NewLine;

                foreach (Object key in children.Keys)
                {
                    s = s + indent + key + Environment.NewLine;

                    s = s + children[key].ToString("  " + indent);
                }

                if (unknown_value != null)
                {
                    s = s + indent + "Unknown values" + Environment.NewLine;

                    s = s + unknown_value.ToString("  " + indent);

                }
                return s;
            }
        }

        public override String ToString()
        {
            return this.ToString("");
        }
    }
}
